#include <string.h>
#include "transenc.h"
#include "hfile.h"
#include "encflagdata.h"

unsigned char gs_default_key[] =
{
    'S', 'h', 'a', 'n', 'd', 'o', 'n', 'g',
    'G', 'o', 'l', 'd', 'e', 'n', 'c', 'i',
    's', 's', 'o', 'f', 't', 'w', 'a', 'r',
    'e', 'T', 'r', 'a', 'n', 's', 'e', 'n',
    '\0'
};

static int _check_intersection(int* keys1, int count1, int* keys2, int count2)
{
    int i = 0;
    int j = 0;

    for (i = 0; i < count1; ++i)
    {
        for (j = 0; j < count2; ++j)
        {
            if (keys1[i] == keys2[j])
            {
                return 0;
            }
        }
    }

    return 1;
}

static int _check_enc_keys(int* keys, int key_count)
{
    int idx = 0;
    if (key_count > DOCTE_ENC_MAXKEYCOUNT)
    {
        return 2;
    }

    for (idx = 0; idx < key_count; ++idx)
    {
        if (!keys[idx])
        {
            return 1;
        }
    }

    return 0;
}

int HPLATFORMCALL _check_enced_file(hfile_t fd, int* keys, int* key_count, char* aes_key)
{
    char buff[ENC_FLAG_BYTES] = {0};
    PENC_FLAG_DATA hdr = (PENC_FLAG_DATA)buff;
    hfile_off_t file_size = hfile_getsize(fd);
    if (file_size < ENC_FLAG_BYTES)
    {
        return 1;
    }

    hfile_offset(fd, 0, HOM_BEGIN);
    if (ENC_FLAG_BYTES != hfile_read(fd, buff, ENC_FLAG_BYTES))
    {
        return 1;
    }

    if (0 != strcmp(hdr->strFlag, IDENTITY_FLAG))
    {
        return 1;
    }

    if (keys && key_count)
    {
        int idx = 0;
        int size = *key_count;
        *key_count = 0;

        for (idx = 0; idx < size && idx < DOCTE_ENC_MAXKEYCOUNT; ++idx)
        {
            if (!hdr->encKeys[idx])
            {
                break;
            }

            keys[idx] = hdr->encKeys[idx];
            (*key_count)++;
        }
    }

    if (aes_key)
    {
        memcpy(aes_key, hdr->realKey, sizeof(hdr->realKey));
    }

    return 0;
}

static int _docte_encrypt_to(hfile_t src, hfile_t dst, int* keys, int key_count)
{
    int ret = 1;

    do
    {
        char hdr_buff[ENC_FLAG_BYTES] = {0};
        char buff[4096] = {0};
        PENC_FLAG_DATA hdr = (PENC_FLAG_DATA)hdr_buff;
        int buff_size = 0;

        if (0 == _check_enced_file(src, NULL, NULL, NULL))
        {
            break;
        }
        if (_check_enc_keys(keys, key_count))
        {
            break;
        }

        hfile_offset(dst, 0, HOM_BEGIN);
        hfile_offset(src, 0, HOM_BEGIN);

        memcpy(hdr->strFlag, IDENTITY_FLAG, sizeof(IDENTITY_FLAG));
        hdr->nVerNum = FILE_HEADER_VER;
        memcpy(hdr->realKey, gs_default_key, sizeof(hdr->realKey));
        if (key_count)
        {
            int i = 0;
            for (i = 0; i < key_count; ++i)
            {
                hdr->encKeys[i] = keys[i];
            }
        }
        hfile_write(dst, (const char*)hdr, ENC_FLAG_BYTES);

        InitAesKey(gs_default_key, CRYPTO_KEY_CHARS);
        ret = 0;
        while ((buff_size = hfile_read(src, buff, 4096)) > 0)
        {
            AesEnc(buff, buff_size, AES_ENC_FLAG);
            if (buff_size != hfile_write(dst, buff, buff_size))
            {
                ret = 1;
                break;
            }
        }
    } while (0);

    return ret;
}

static int _docte_decrypt_to(hfile_t src, hfile_t dst, int* keys, int key_count)
{
    int ret = 1;

    do
    {
        char aes_key[CRYPTO_KEY_CHARS] = {0};
        char buff[4096];
        int buff_size = 0;
        int file_keys[DOCTE_ENC_MAXKEYCOUNT] = {0};
        int file_key_count = DOCTE_ENC_MAXKEYCOUNT;
        if (0 != _check_enced_file(src, file_keys, &file_key_count, aes_key))
        {
            break;
        }

        if (file_key_count && _check_intersection(keys, key_count, file_keys, file_key_count))
        {
            break;
        }

        InitAesKey((unsigned char*)aes_key, CRYPTO_KEY_CHARS);

        ret = 0;
        while ((buff_size = hfile_read(src, buff, 4096)) > 0)
        {
            AesDec(buff, buff_size, AES_ENC_FLAG);
            if (buff_size != hfile_write(dst, buff, buff_size))
            {
                ret = 1;
                break;
            }
        }
    } while (0);

    return ret;
}

int HPLATFORMCALL docte_encrypt_to(const char* enc_this, const char* to_this, int* keys, int key_count)
{
    int ret = 1;
    hfile_t src = HF_INVALID_FILE;
    hfile_t dst = HF_INVALID_FILE;

    do
    {
        if (!enc_this || !to_this) // support keys later
        {
            break;
        }

        src = hfile_open(enc_this, HOA_READ, HOF_OPEN);
        dst = hfile_open(to_this, HOA_WRITE, HOF_CREATE);
        if (HF_INVALID_FILE == src || HF_INVALID_FILE == dst)
        {
            break;
        }

        ret = _docte_encrypt_to(src, dst, keys, key_count);
    } while (0);

    if (HF_INVALID_FILE != src)
    {
        hfile_close(src);
    }

    if (HF_INVALID_FILE != dst)
    {
        hfile_close(dst);
    }

    return ret;
}

#if _MSC_VER
int HPLATFORMCALL docte_encrypt_to16(const wchar_t* enc_this, const wchar_t* to_this, int* keys, int key_count)
{
    int ret = 1;
    hfile_t src = HF_INVALID_FILE;
    hfile_t dst = HF_INVALID_FILE;

    do
    {
        if (!enc_this || !to_this) // support keys later
        {
            break;
        }

        src = hfile_open16(enc_this, HOA_READ, HOF_OPEN);
        dst = hfile_open16(to_this, HOA_WRITE, HOF_CREATE);
        if (HF_INVALID_FILE == src || HF_INVALID_FILE == dst)
        {
            break;
        }

        ret = _docte_encrypt_to(src, dst, keys, key_count);
    } while (0);

    if (HF_INVALID_FILE != src)
    {
        hfile_close(src);
    }

    if (HF_INVALID_FILE != dst)
    {
        hfile_close(dst);
    }

    return ret;
}
#endif // #if _MSC_VER

int HPLATFORMCALL docte_decrypt_to(const char* dec_this, const char* to_this, int* keys, int key_count)
{
    int ret = 1;
    hfile_t src = HF_INVALID_FILE;
    hfile_t dst = HF_INVALID_FILE;

    do
    {
        if (!dec_this || !to_this) // support keys later
        {
            break;
        }

        src = hfile_open(dec_this, HOA_READ, HOF_OPEN);
        dst = hfile_open(to_this, HOA_WRITE, HOF_CREATE);
        if (HF_INVALID_FILE == src || HF_INVALID_FILE == dst)
        {
            break;
        }

        ret = _docte_decrypt_to(src, dst, keys, key_count);
    } while (0);

    if (HF_INVALID_FILE != src)
    {
        hfile_close(src);
    }

    if (HF_INVALID_FILE != dst)
    {
        hfile_close(dst);
    }

    return ret;
}

#if _MSC_VER
int HPLATFORMCALL docte_decrypt_to16(const wchar_t* dec_this, const wchar_t* to_this, int* keys, int key_count)
{
    int ret = 1;
    hfile_t src = HF_INVALID_FILE;
    hfile_t dst = HF_INVALID_FILE;

    do
    {
        if (!dec_this || !to_this) // support keys later
        {
            break;
        }

        src = hfile_open16(dec_this, HOA_READ, HOF_OPEN);
        dst = hfile_open16(to_this, HOA_WRITE, HOF_CREATE);
        if (HF_INVALID_FILE == src || HF_INVALID_FILE == dst)
        {
            break;
        }

        ret = _docte_decrypt_to(src, dst, keys, key_count);
    } while (0);

    if (HF_INVALID_FILE != src)
    {
        hfile_close(src);
    }

    if (HF_INVALID_FILE != dst)
    {
        hfile_close(dst);
    }

    return ret;
}
#endif // #if _MSC_VER

int HPLATFORMCALL docte_getenckeys(const char* file_name, int* keys, int* key_count)
{
    int ret = 1;

    if (file_name)
    {
        hfile_t fd = hfile_open(file_name, HOA_READ, HOF_OPEN);
        if (HF_INVALID_FILE != fd)
        {
            ret = _check_enced_file(fd, keys, key_count, NULL);
            hfile_close(fd);
        }
    }

    return ret;
}

#if _MSC_VER
int HPLATFORMCALL docte_getenckeys16(const wchar_t* file_name, int* keys, int* key_count)
{
    int ret = 1;

    if (file_name)
    {
        hfile_t fd = hfile_open16(file_name, HOA_READ, HOF_OPEN);
        if (HF_INVALID_FILE != fd)
        {
            ret = _check_enced_file(fd, keys, key_count, NULL);
            hfile_close(fd);
        }
    }

    return ret;
}
#endif // #if _MSC_VER

static int _setenckeys(hfile_t fp, const int* keys, int key_count)
{
    int ret = 1;
    char buff[ENC_FLAG_BYTES] = {0};
    PENC_FLAG_DATA hdr = (PENC_FLAG_DATA)buff;
    hfile_off_t file_size = 0;
    int enc_count = 0;
    do 
    {
        if (HF_INVALID_FILE == fp)
        {
            break;
        }

        if (0 != _check_enc_keys((int*)keys, key_count))
        {
            break;
        }

        file_size = hfile_getsize(fp);
        if (file_size < ENC_FLAG_BYTES)
        {
            break;
        }

        hfile_offset(fp, 0, HOM_BEGIN);
        if (ENC_FLAG_BYTES != hfile_read(fp, buff, ENC_FLAG_BYTES))
        {
            break;
        }

        if (0 != strcmp(hdr->strFlag, IDENTITY_FLAG))
        {
            break;
        }

        memset(hdr->encKeys, 0x00, sizeof(hdr->encKeys));
        if (keys && key_count)
        {
            int idx = 0;
            for (idx = 0; idx < key_count && idx < DOCTE_ENC_MAXKEYCOUNT; idx++)
            {
                hdr->encKeys[idx] = keys[idx];
            }
        }
        hfile_offset(fp, 0, HOM_BEGIN);
        if (sizeof(ENC_FLAG_DATA) == hfile_write(fp, (const char *)hdr, sizeof(ENC_FLAG_DATA)))
        {
            ret = 0;
        }
    } while (0);
    return ret;
}

//
int HPLATFORMCALL docte_setenckeys(const char *file_name, const int* keys, int key_count)
{
    int ret = 1;
    hfile_t fp = HF_INVALID_FILE;
    if (!file_name || !file_name[0])
    {
        return ret;
    }
    fp = hfile_open(file_name, HOA_READ | HOA_WRITE, HOF_OPEN);
    if (HF_INVALID_FILE == fp)
    {
        return ret;
    }
    ret = _setenckeys(fp, keys, key_count);
    hfile_close(fp);
    return ret;
}

#if _MSC_VER
int HPLATFORMCALL docte_setenckeys16(const wchar_t* file_name, const int* keys, int key_count)
{
    int ret = 1;
    hfile_t fp = HF_INVALID_FILE;
    if (!file_name || !file_name[0])
    {
        return ret;
    }
    fp = hfile_open16(file_name, HOA_READ | HOA_WRITE, HOF_OPEN);
    if (HF_INVALID_FILE == fp)
    {
        return ret;
    }
    ret = _setenckeys(fp, keys, key_count);
    hfile_close(fp);
    return ret;
}
#endif // #if _MSC_VER

int HPLATFORMCALL docte_decrypt_test(const char* file_name, int* keys, int key_count)
{
    int enc_keys[DOCTE_ENC_MAXKEYCOUNT];
    int enc_key_count = DOCTE_ENC_MAXKEYCOUNT;
    if (0 != docte_getenckeys(file_name, enc_keys, &enc_key_count))
    {
        return 1;
    }

    // if encrypted without key, the file can decrypt by any key
    if (!enc_key_count)
    {
        return 0;
    }

    return _check_intersection(keys, key_count, enc_keys, enc_key_count);
}

#if _MSC_VER
int HPLATFORMCALL docte_decrypt_test16(const wchar_t* file_name, int* keys, int key_count)
{
    int enc_keys[DOCTE_ENC_MAXKEYCOUNT];
    int enc_key_count = DOCTE_ENC_MAXKEYCOUNT;
    if (0 != docte_getenckeys16(file_name, enc_keys, &enc_key_count))
    {
        return 1;
    }

    if (!enc_key_count)
    {
        return 0;
    }

    return _check_intersection(keys, key_count, enc_keys, enc_key_count);
}
#endif // #if _MSC_VER
